import json
import streamlit as st
import requests
import sseclient

st.set_page_config(page_title="LangGraph Queue SSE Demo", page_icon="⚡")
st.title("LangGraph 实时进度展示 (Queue + SSE)")

st.markdown("点击下方按钮开始执行流程，并实时接收中间进度、节点结果以及完成信号。")

if st.button("开始执行流程"):
    # 页面占位符
    progress_box = st.empty()
    graph_box = st.empty()
    status_box = st.empty()
    log_box = st.container()
    log_box.write("等待数据中...")

    # 与 FastAPI 后端建立 SSE 连接
    with requests.get("http://localhost:8002/stream", stream=True) as r:
        client = sseclient.SSEClient(r)
        for msg in client.events():
            if not msg.data.strip():
                continue
            try:
                data = json.loads(msg.data)
            except json.JSONDecodeError:
                log_box.write(f"⚠️ 无法解析: {msg.data}")
                continue

            msg_type = data.get("type")

            # 🔄 进度事件
            if msg_type == "progress":
                loop_idx = data["loop"]
                progress_box.write(f"🔄 循环进度：第 {loop_idx + 1} 次")
                log_box.write(f"Progress Event: {data}")

            # ✅ LangGraph节点事件
            elif msg_type == "graph":
                node = data["node"]
                graph_box.write(f"✅ 节点 [{node}] 执行完成: {json.dumps(data['state'], ensure_ascii=False)}")
                log_box.write(f"Graph Event: {data}")

            # 🎯 完成事件
            elif msg_type == "DONE":
                status_box.success("🎯 流程执行完毕！")
                log_box.write("DONE Event: 流程结束")
                break
